// @HEADER
//*********************************************************************//
//  SiYuan: A numerical PDE solver                                     //
//  Copyright (2022) YUAN Xi                                           //
//  This Software is released under the BSD 2-Clause license detailed  //
//  in the file "LICENSE" in the top-level SiYuan directory            //
//*********************************************************************//
// @HEADER

#ifndef _GradShafranov_EquationSet_impl_hpp_
#define _GradShafranov_EquationSet_impl_hpp_

#include "Teuchos_ParameterList.hpp"
#include "Teuchos_StandardParameterEntryValidators.hpp"
#include "Teuchos_RCP.hpp"
#include "Teuchos_Assert.hpp"
#include "Phalanx_DataLayout_MDALayout.hpp"
#include "Phalanx_FieldManager.hpp"

#include "Panzer_IntegrationRule.hpp"
#include "Panzer_BasisIRLayout.hpp"

// include evaluators here
#include "Panzer_Integrator_BasisTimesScalar.hpp"
#include "Panzer_Integrator_TransientBasisTimesScalar.hpp"
#include "Panzer_Integrator_BasisTimesVector.hpp"
#include "Panzer_Integrator_BasisTimesTensorTimesVector.hpp"
#include "Panzer_Product.hpp"

#include <bits/stdc++.h>   // for M_PI

// ***********************************************************************
template <typename EvalT>
SiYuan::EquationSet_GradShafranov<EvalT>::
EquationSet_GradShafranov(const Teuchos::RCP<Teuchos::ParameterList>& params,
    const int& default_integration_order,
    const panzer::CellData& cell_data,
    const Teuchos::RCP<panzer::GlobalData>& global_data,
    const bool build_transient_support) :
    EquationSet<EvalT> (params,default_integration_order,cell_data,global_data,build_transient_support ) 
{
  std::string model_id = params->get<std::string>("Model ID");
  std::string basis_type = "HGrad";
  int basis_order = params->get<int>("Basis Order");
  int integration_order = params->get<int>("Integration Order");

  if (!this->buildTransientSupport()) {
    m_dof_name = "Magnetic_Flux";  //  magnetic flux function

    this->addDOF(m_dof_name,basis_type,basis_order,integration_order);
    this->addDOFGrad(m_dof_name);
  } else {
  /*  m_dof_name = "Magnetic_Flux";
    this->enable_xdotdot();

    this->addDOF(m_dof_names[0],basis_type,basis_order,integration_order);
    this->addDOFGrad(m_dof_names[0]);
    this->addDOFTimeDerivative(m_dof_names[0]);
    this->addDOFDotDot(m_dof_names[0]);*/
  }

  this->addClosureModel(model_id);
  this->setupDOFs();
}

// ***********************************************************************
template <typename EvalT>
void SiYuan::EquationSet_GradShafranov<EvalT>::
buildAndRegisterEquationSetEvaluators(PHX::FieldManager<panzer::Traits>& fm,
    const panzer::FieldLibrary& /* fl */,
    const Teuchos::ParameterList& ) const
{
  using panzer::EvaluatorStyle;

  Teuchos::RCP<panzer::IntegrationRule> ir = this->getIntRuleForDOF(m_dof_name); 
  Teuchos::RCP<panzer::BasisIRLayout> basis = this->getBasisIRLayoutForDOF(m_dof_name);

  // Residual
    Teuchos::RCP<PHX::Evaluator<panzer::Traits>> top =Teuchos:: rcp(new
      SiYuan::Residual_GradShafranov<EvalT, panzer::Traits>(EvaluatorStyle::EVALUATES,
        *basis, *ir ) );
    this->template registerEvaluator<EvalT>(fm, top);
}


// ***********************************************************************
/////////////////////////////////////////////////////////////////////////////
//
//  Main Constructor: Residual_Helmholtz
//
/////////////////////////////////////////////////////////////////////////////
template<typename EvalT, typename Traits>
SiYuan::Residual_GradShafranov<EvalT, Traits>::
Residual_GradShafranov(
    const panzer::EvaluatorStyle&   evalStyle,
    const panzer::BasisIRLayout&    basis,
    const panzer::IntegrationRule&  ir )
    : evalStyle_(evalStyle), basisName_(basis.name())
{
    irDegree_ = ir.cubature_degree;
    mgflux_ = PHX::MDField<const ScalarT, panzer::Cell, panzer::IP>("Magnetic_Flux", ir.dl_scalar);
  //  this->addDependentField(mgflux_);

    grad_mgflux_ = PHX::MDField<const ScalarT, panzer::Cell, panzer::IP, panzer::Dim>("GRAD_Magnetic_Flux", ir.dl_vector);
    this->addDependentField(grad_mgflux_);

    
    // Create the field that we're either contributing to or evaluating
    // (storing).
    residual_ = PHX::MDField<ScalarT, panzer::Cell, panzer::BASIS>("RESIDUAL_Magnetic_Flux",basis.functional);
    if (evalStyle == panzer::EvaluatorStyle::CONTRIBUTES) 
    {
      this->addContributedField(residual_);
    }
    else // if (evalStyle == EvaluatorStyle::EVALUATES)
    {
      this->addEvaluatedField(residual_);
    }

    // Set the name of this object.
    std::string n("MagneticFlux (");
    if (evalStyle == panzer::EvaluatorStyle::CONTRIBUTES)
      n += "CONTRIBUTES)";
    else // if (evalStyle == EvaluatorStyle::EVALUATES)
      n += "EVALUATES)";
  //  n += "):  " + field_.fieldTag().name();
    this->setName(n);
} // end of Main Constructor

/////////////////////////////////////////////////////////////////////////////
//
//  postRegistrationSetup()
//
/////////////////////////////////////////////////////////////////////////////
template<typename EvalT, typename Traits>
void
SiYuan::Residual_GradShafranov<EvalT, Traits>:: postRegistrationSetup(
    typename Traits::SetupData sd, PHX::FieldManager<Traits>& fm)
{
    basisIndex_ = panzer::getBasisIndex(basisName_, (*sd.worksets_)[0], this->wda);
    irIndex_ = panzer::getIntegrationRuleIndex(irDegree_,(*sd.worksets_)[0], this->wda);
} // end of postRegistrationSetup()

/////////////////////////////////////////////////////////////////////////////
//
//  evaluateFields()
//
/////////////////////////////////////////////////////////////////////////////
template<typename EvalT, typename Traits>
void
SiYuan::Residual_GradShafranov<EvalT, Traits>::
evaluateFields( typename Traits::EvalData workset )
{
    using Kokkos::parallel_for;
    using Kokkos::RangePolicy;
    using PHX::Device;

    typedef Intrepid2::FunctionSpaceTools<PHX::exec_space> FST;

//std::cout << "RES:" << PHX::print<EvalT>()  << std::endl;
    // Grab the basis information.
    const panzer::BasisValues2<double>& bv = *this->wda(workset).bases[basisIndex_];
    auto ip_coordinates = workset.int_rules[irIndex_]->ip_coordinates.get_static_view();

    PHX::MDField<double, panzer::Cell, panzer::BASIS, panzer::IP, panzer::Dim>
          weightedGradBasis = bv.weighted_grad_basis;
    PHX::MDField<double, panzer::Cell, panzer::BASIS, panzer::IP>
          weightedBasis = bv.weighted_basis_scalar;

    int numDims = weightedGradBasis.extent(3);
    int numBases= weightedGradBasis.extent(1);
    int numQP = weightedGradBasis.extent(2);

    for( int cell=0; cell<workset.num_cells; cell++ )
    {
        for (int basis(0); basis < numBases; ++basis)
        {
          if (evalStyle_ == panzer::EvaluatorStyle::EVALUATES) {
            residual_(cell, basis) = 0.0;
          }
          for (int qp(0); qp < numQP; ++qp) {
            for (int dim(0); dim < numDims; ++dim) {
              residual_(cell, basis) += grad_mgflux_(cell, qp, dim) * weightedGradBasis(cell, basis, qp, dim);
            }
            const double & r = ip_coordinates(cell,qp,0);
            residual_(cell, basis) += grad_mgflux_(cell, qp, 0) * weightedBasis(cell, basis, qp)/r;
            residual_(cell, basis) -= (r*r+1.0) * weightedBasis(cell, basis, qp);
          }
        }
    }

  //  std::cout << h_(0,0) << "," << workset.num_cells << std::endl;

} // end of evaluateFields()


#endif 
